#include <vector>
#include <iostream>
#include "gtest/gtest.h"
#include <zookeeper.h>
#include "glog/logging.h"
#include "fmt/format.h"
#include <thread>
#include <chrono>
#include <mutex>
#include <condition_variable>
class Status{
public:
    enum ErrorCode {
        SUCCESS = 0,
        FAIL = 1,
    };
    Status(){}
    Status(int c = 0):code(c) {}
    template<typename ...Args>
    Status(Args ...args) : code(FAIL), msg(fmt::format(std::forward<Args>(args)...)) {
    }
    template<typename ...Args>
    Status(int c, Args ...args) : code(c), msg(fmt::format(std::forward<Args>(args)...)) {
 
    }
    operator bool() const {
        return code == SUCCESS;
    }
    int code = 0;
    std::string msg;
    ~Status(){
        if(code != SUCCESS) {
            LOG(INFO) << "code = " << code << " msg = " << msg;
        }
    }
};
 
#define CHECK_COND_RETURN(cond, ...) \
do{\
    if(!(cond)){\
        return Status(__VA_ARGS__);\
    }\
}while(0)
 
class ZooLock {
public:
    ZooLock(const std::string &addr, const std::string &path)
        : zk_addr(addr), lock_path(path)
    {
        
    }
    Status Init() {
        handle = zookeeper_init(zk_addr.c_str(), watcher_fn, recv_timeout, nullptr, this, 0);
        CHECK_COND_RETURN(handle != nullptr, "zk init fail, return null handle");
        LOG(INFO) << "zk init ok, handle is: " << handle;
        return 0;
    }
    Status Lock() {
        // int zoo_create(zhandle_t *zh, const char *path, const char *value, int valuelen,
        //        const struct ACL_vector *acl, int flags, char *path_buffer,
        //        int path_buffer_len)
        //create father path
        auto ret = zoo_create(
            handle, "/lock",
            nullptr, 0,
            &ZOO_OPEN_ACL_UNSAFE,
            0,
            nullptr, 0);
        CHECK_COND_RETURN(
            ret == ZOK || ret == ZNODEEXISTS,
            "zk create /lock fail, return null handle");
 
        //create tmp and seq node
        char path_buff[256] = {0};
        CHECK_COND_RETURN(
            zoo_create(
                handle, lock_path.c_str(),
                nullptr, 0,
                &ZOO_OPEN_ACL_UNSAFE,
                ZOO_EPHEMERAL | ZOO_SEQUENCE,
                path_buff, 256) == ZOK,
            "zk create fail, return null handle");
        /*int zoo_get_children(zhandle_t *zh, const char *path, int watch,
                     struct String_vector *strings)
        */
        self_node_path = std::string(path_buff);
        self_node_path = self_node_path.substr(6, self_node_path.size());
        //check self is min?
        do{
            struct String_vector strings;
            CHECK_COND_RETURN(
                zoo_get_children(handle, "/lock", 0, &strings) == ZOK,
                "zk get children fail, return null handle");
            std::vector<std::string> all_path;
            for(int i = 0; i < strings.count; i++) {
                all_path.emplace_back(strings.data[i]);
                delete strings.data[i];
            }
            std::sort(all_path.begin(), all_path.end());
            // for(auto &path : all_path) {
            //     LOG(INFO) << "all path: " << path;
            // }
            CHECK_COND_RETURN(all_path.size() > 0, "no path, create failed ?");
            if(all_path[0] == self_node_path) { //self is min, so get the lock
                break;
            }
            bool found = false;
            for(int i = 0; i < all_path.size(); i++) {
                if(all_path[i] ==  self_node_path) {
                    found = true;
                    min_node_path = "/lock/" + all_path[i-1];
                    break;
                }
            }
            // zoo_get(zhandle_t *zh, const char *path, int watch, char *buffer,
            // int *buffer_len, struct Stat *stat)
            struct Stat stat;
            int len = 256;
            //get min node, add watch.  int watc = 1, will call watcher_fn
            auto ret = zoo_get(handle, min_node_path.c_str(), 1, path_buff, &len, &stat);
            if(ret == ZNONODE){
                continue;
            }
            CHECK_COND_RETURN(ret == ZOK, "add watch to {}, failed code: {}", min_node_path, ret);
            CHECK_COND_RETURN(found, "no path found, create failed ?");
            std::unique_lock<std::mutex> guard(_mtx);
            cv.wait(guard, [this](){return min_node_deleted;});
        }while(true);
        LOG(INFO) << "create success: " << path_buff;
        return 0;
    }
    Status UnLock() {
        auto path = "/lock/" + self_node_path;
        CHECK_COND_RETURN(zoo_delete(handle, path.c_str() , -1) == ZOK, "close zk fail");
        // handle = nullptr;
        return 0;
    }
 
    static void watcher_fn(zhandle_t *zh, int type,
        int state, const char *path,void *watcherCtx)
    {
        ZooLock *self = (ZooLock*) watcherCtx;
        LOG(INFO) << "zk state, self " << self 
                  << " type: " << type
                  << " state: " << state 
                  << " path: " << path
                  << " min_node_path: " << self->min_node_path;
        if(path == self->min_node_path) {
            if(type == ZOO_DELETED_EVENT) { //node deleted, so awake
                LOG(INFO) << "min node deleted";
                std::unique_lock<std::mutex> guard(self->_mtx);
                self->min_node_deleted = true;
                self->cv.notify_all();
            }
        }
    }
    ~ZooLock() {
        zookeeper_close(handle);
    }
private:
    std::string zk_addr;
    std::string lock_path;
    zhandle_t *handle = nullptr;
    int recv_timeout = 5000;
    std::string self_node_path;
    std::condition_variable cv;
    std::mutex _mtx;
    bool min_node_deleted = false;
    std::string min_node_path;
};
 
 
TEST(zk, zk) {
    {
        std::vector<std::thread> ths;
        int counter = 0;
        int target = 20;
        int num = 30;
        for(int i = 0; i < target; i++) {
            ths.emplace_back([i, &counter, num](){
                ZooLock lock2("localhost:2181", "/lock/ckpt-");
                 ASSERT_TRUE(lock2.Init());
                for(int j = 0; j < num; j++) {
                    ASSERT_TRUE(lock2.Lock());
                    auto tmp = counter;
                    tmp++;
                    // std::this_thread::sleep_for(std::chrono::milliseconds(std::rand() % 1000));
                    counter = tmp;
                    LOG(INFO) << "after lock got " << i << " counter: " << counter;
                    ASSERT_TRUE(lock2.UnLock());
                }
            });
        }
        for(auto &t : ths) {
            t.join();
        }
        ASSERT_EQ(counter, target * num);
    }
}